import markitdown
import pypandoc
import os
from docx import Document
from docx.shared import Pt, RGBColor
from docx.oxml.ns import qn
from docx.enum.style import WD_STYLE_TYPE
from docx.enum.table import WD_TABLE_ALIGNMENT

def create_template_docx(template_path="template.docx"):
    doc = Document()
    
    # 定义要设置的所有样式名称
    style_names = [
        "Normal", "Heading 1", "Heading 2", "Heading 3", "Heading 4",
        "List Bullet", "List Number", "Table Grid", "Caption",
        "Body Text", "Quote", "Intense Quote", "Subtitle",
        "List Paragraph", "TOC 1", "TOC 2", "TOC 3"
    ]
    
    # 为所有样式设置默认格式
    for style_name in style_names:
        if style_name not in doc.styles:
            if "TOC" in style_name:
                style = doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
            elif style_name == "Table Grid":
                style = doc.styles.add_style(style_name, WD_STYLE_TYPE.TABLE)
            else:
                try:
                    style = doc.styles[style_name]
                except:
                    style = doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
        else:
            style = doc.styles[style_name]
            
        # 设置基本字体属性
        if hasattr(style, 'font'):
            style.font.name = "Times New Roman"
            style._element.rPr.rFonts.set(qn("w:eastAsia"), "宋体")
            style.font.color.rgb = RGBColor(0, 0, 0)
            
            # 根据样式类型设置不同的字号和粗细
            if "Heading 1" in style_name:
                style.font.size = Pt(16)
                style.font.bold = True
            elif "Heading 2" in style_name:
                style.font.size = Pt(14)
                style.font.bold = True
            elif "Heading 3" in style_name:
                style.font.size = Pt(13)
                style.font.bold = True
            elif "Heading 4" in style_name:
                style.font.size = Pt(12)
                style.font.bold = True
            else:
                style.font.size = Pt(12)
                style.font.bold = False
    
    # 添加示例表格以设置默认网格线
    table = doc.add_table(rows=1, cols=1)
    table.style = "Table Grid"
    table.alignment = WD_TABLE_ALIGNMENT.CENTER
    
    # 设置表格样式
    for row in table.rows:
        for cell in row.cells:
            cell.paragraphs[0].add_run("Sample")
            for border in ["top", "left", "bottom", "right"]:
                tc = cell._tc
                tc.get_or_add_tcPr()
                tcBorders = tc.get_or_add_tcPr().get_or_add_tcBorders()
                border_element = getattr(tcBorders, f"get_or_add_{border}")()
                border_element.val = "single"
                border_element.sz = 4
                border_element.color = "000000"
    
    doc.save(template_path)
    print(f"Template saved to {template_path}")



# 提取任意类型文件-》.md
def extract_markdown(input_file, output_md="output.md"):
    """
    Extract Markdown from input file using MarkItDown.
    input_file: Path to input file (e.g., PDF, MD).
    output_md: Path to save extracted Markdown.
    """
    try:
        # Use MarkItDown to extract Markdown
        result = markitdown.convert(input_file)
        # Save the extracted Markdown to a file
        with open(output_md, "w", encoding="utf-8") as f:
            f.write(result)
        print(f"Markdown extracted to {output_md}")
        return output_md
    except Exception as e:
        print(f"Error extracting Markdown: {e}")
        return None

# 将.md转换为.docx
def convert_md_to_docx_pure(input_md, output_docx="output.docx"):
    """
    Convert Markdown to DOCX using pypandoc.
    input_md: Path to Markdown file.
    output_docx: Path to save DOCX file.
    """
    try:
        # Convert Markdown to DOCX
        pypandoc.convert_file(
            input_md,
            "docx",
            outputfile=output_docx,
            extra_args=["--standalone", "--toc"]  # Optional: Add table of contents
        )
        print(f"DOCX saved to {output_docx}")
    except Exception as e:
        print(f"Error converting to DOCX: {e}")

def convert_md_to_docx_template(input_md, output_docx="output.docx", template_path="template.docx"):
    """Convert Markdown to DOCX using pypandoc with a reference template."""
    try:
        # 确保模板存在
        if not os.path.exists(template_path):
            print(f"Template not found at {template_path}, creating new template...")
            create_template_docx(template_path)
        
        # 转换Markdown到DOCX
        pypandoc.convert_file(
            input_md,
            "docx",
            outputfile=output_docx,
            extra_args=[
                "--reference-doc", template_path,
                "--toc",  # 添加目录
                "-f", "markdown+pipe_tables+fenced_code_blocks+east_asian_line_breaks",  # 支持表格、代码块和中文换行
                "--wrap=none",  # 防止不必要的换行
                "--toc-depth=3",  # 目录深度设置为3级
                "--columns=10000"  # 防止内容被过早换行
            ]
        )
        print(f"DOCX saved to {output_docx}")
    except Exception as e:
        print(f"Error converting to DOCX: {e}")
        
def main():
    # # Input file (e.g., PDF or existing Markdown)
    # input_file = "input.pdf"  # Replace with your file path
    # output_md = "output.md"
    # output_docx = "output.docx"

    
    # # Step 1: Extract Markdown
    # md_file = extract_markdown(input_file, output_md)
    # if not md_file or not os.path.exists(md_file):
    #     print("Markdown extraction failed.")
    #     return
    # # Run to create template
    # create_template_docx()
    
    # Step 2: Convert to DOCX
    md_file="input_files/output_image_寒武纪2023合并资产负债表-图片式.md"
    output_docx="output_files/output_image_寒武纪2023合并资产负债表-图片式.docx"
    template_path="template.docx"
    convert_md_to_docx_template(md_file, output_docx,template_path)

if __name__ == "__main__":
    main()